package com.lw.leetcode.hash.b;

import java.util.*;

/**
 * Created with IntelliJ IDEA.
 * 2502. 设计内存分配器
 *
 * @author liw
 * @version 1.0
 * @date 2022/12/13 11:23
 */
public class Allocator {

    public static void main(String[] args) {

        int n = 10;
        Allocator test = new Allocator(n);

        // [null, 0, 1, 2, 1, 3, 1, 6, 3, -1, 0]
        System.out.println(test.allocate(1, 1));
        System.out.println(test.allocate(1, 2));
        System.out.println(test.allocate(1, 3));
        System.out.println(test.free(2));
        System.out.println(test.allocate(3, 4));
        System.out.println(test.allocate(1, 1));
        System.out.println(test.allocate(1, 1));
        System.out.println(test.free(1));
        System.out.println(test.allocate(10, 2));
        System.out.println(test.free(7));


    }

    private int n;
    private Map<Integer, List<Integer>> map = new HashMap<>();
    private TreeMap<Integer, Integer> treeMap = new TreeMap<>();

    public Allocator(int n) {
        this.n = n;
    }

    public int allocate(int size, int mID) {
        int st = 0;
        for (Map.Entry<Integer, Integer> entry : treeMap.entrySet()) {
            Integer key = entry.getKey();
            Integer value = entry.getValue();
            if (key - st >= size) {
                treeMap.put(st, st + size - 1);
                map.computeIfAbsent(mID, v -> new ArrayList<>()).add(st);
                return st;
            } else {
                st = value + 1;
            }
        }
        if (st + size > n) {
            return -1;
        }
        treeMap.put(st, st + size - 1);
        map.computeIfAbsent(mID, v -> new ArrayList<>()).add(st);
        return st;
    }

    public int free(int mID) {
        List<Integer> list = map.get(mID);
        if (list == null) {
            return 0;
        }
        int sum = 0;
        for (int index : list) {
            sum += (treeMap.remove(index) - index + 1);
        }
        map.remove(mID);
        return sum;
    }

}
